#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from pymongo import MongoClient
from datetime import datetime
import fire


# 连接MongoDB
def InitMongoDB(MONGODB_HOST, MONGODB_PORT, MONGODB_DB_NAME, MONGODB_COLLECTION_NAME):
    client = MongoClient(MONGODB_HOST, MONGODB_PORT)
    db = client[MONGODB_DB_NAME]
    collection = db[MONGODB_COLLECTION_NAME]
    return client, collection


# 关闭MongoDB
def CloseMongoDB(MongoDB_Connection):
    MongoDB_Connection.close()


def TypeRawdata2Datalist(type: str):
    client_raw, collection_raw = InitMongoDB("10.10.116.28", 27017, "PD_Data", "rawdata")
    client_antnest, collection_antnest = InitMongoDB("10.10.116.28", 27017, "ant_nest", "datalist")

    typedict = {}
    typedict["Gut"] = ["microbiome", "human gut"]
    typedict["cancer"] = ["precision medicine", "cancer"]
    typedict["human and animal gut microbiome"] = ["microbiome", "human and animal gut"]
    typedict["gut microbiome"] = ["microbiome", "human gut"]
    typedict["Host-associated:Human:Respiratory system:Pulmonary system:Lung "] = ["microbiome", "human lung"]
    typedict["hydrosphere"] = ["microbiome", "hydrosphere"]

    for result in collection_raw.find(
            {"type": type, "Yao_to_downloader": {"$exists": 0}, "data_store_path": {"$exists": 0}}):
        record = {}
        if "BioprojectID" in result:
            record["BioprojectID"] = result["BioprojectID"]
        if "Study" in result:
            record["StudyID"] = result["Study"]
        if "RunID" in result:
            record["RunID"] = result["RunID"]
            data_download_url = "anonftp@ftp.ncbi.nlm.nih.gov:/sra/sra-instant/reads/ByRun/sra/" + result["RunID"][
                                                                                                   :3] + "/" + result[
                                                                                                                   "RunID"][
                                                                                                               :6] + "/" + \
                                result["RunID"] + "/" + result["RunID"] + ".sra"
            record["data_download_url"] = data_download_url
        if "source" in result:
            record["source"] = result["source"]
        if "backup" in result:
            record["backup"] = result["backup"]
        if "type" in result:
            record["type"] = typedict[result["type"]]
        if "Date_to_Yao" in result:
            record["data_create_date"] = datetime.strptime(result["Date_to_Yao"], '%Y-%m-%d')

        record["download_count"] = 0
        record["NODE"] = "N"

        collection_antnest.insert_one(record)

    CloseMongoDB(client_raw)
    CloseMongoDB(client_antnest)
    print(type + " type records transferred.")


def NoTypeRawdata2Datalist():
    client_raw, collection_raw = InitMongoDB("10.10.116.28", 27017, "PD_Data", "rawdata")
    client_antnest, collection_antnest = InitMongoDB("10.10.116.28", 27017, "ant_nest", "datalist")

    for result in collection_raw.find(
            {"type": {"$exists": 0}, "Yao_to_downloader": {"$exists": 0}, "data_store_path": {"$exists": 0}}):
        record = {}
        if "BioprojectID" in result:
            record["BioprojectID"] = result["BioprojectID"]
        if "Study" in result:
            record["StudyID"] = result["Study"]
        if "RunID" in result:
            record["RunID"] = result["RunID"]
            data_download_url = "anonftp@ftp.ncbi.nlm.nih.gov:/sra/sra-instant/reads/ByRun/sra/" + result["RunID"][
                                                                                                   :3] + "/" + result[
                                                                                                                   "RunID"][
                                                                                                               :6] + "/" + \
                                result["RunID"] + "/" + result["RunID"] + ".sra"
            record["data_download_url"] = data_download_url
        if "source" in result:
            record["source"] = result["source"]
        if "backup" in result:
            record["backup"] = result["backup"]
        if "Date_to_Yao" in result:
            record["data_create_date"] = datetime.strptime(result["Date_to_Yao"], '%Y-%m-%d')

        record["download_count"] = 0
        record["NODE"] = "N"

        collection_antnest.insert_one(record)

    CloseMongoDB(client_raw)
    CloseMongoDB(client_antnest)
    print("No type records transferred.")


def ImportRestData():
    client_raw, collection_raw = InitMongoDB("localhost", 27017, "PD_Data", "rawdata")
    client_antnest, collection_antnest = InitMongoDB("localhost", 27017, "ant_nest", "datalist")

    ht_job = {}
    ht_date = {}
    ht_date2 = {}
    infile = open("downloader_2.txt", "r")
    for line in infile:
        line = line.rstrip()
        splitline = line.split("\t")
        ht_job[splitline[0]] = splitline[2]
        ht_date[splitline[0]] = datetime.strptime(splitline[3], '%Y-%m-%d')
        ht_date2[splitline[0]] = datetime.strptime(splitline[4], '%Y-%m-%d')
    infile.close()

    # 处理已经分配的
    count = 0
    print("Start processing assigned data.")
    for result in collection_raw.find({"JobID": {"$exists": 1}}):
        insert_record = {}
        insert_record["RunID"] = result["RunID"]
        if "BioprojectID" in result:
            insert_record["BioprojectID"] = result["BioprojectID"]
        if "Study" in result:
            insert_record["StudyID"] = result["Study"]
        if "source" in result:
            insert_record["source"] = result["source"]
        data_download_url = "anonftp@ftp.ncbi.nlm.nih.gov:/sra/sra-instant/reads/ByRun/sra/" + result["RunID"][
                                                                                               :3] + "/" + result[
                                                                                                               "RunID"][
                                                                                                           :6] + "/" + \
                            result["RunID"] + "/" + result["RunID"] + ".sra"
        insert_record["data_download_url"] = data_download_url
        if "backup" in result:
            insert_record["backup"] = result["backup"]
        insert_record["download_count"] = 0
        if "type" in result:
            insert_record["type"] = result["type"]
        insert_record["data_create_date"] = ht_date[result["JobID"]]

        insert_record["download_way"] = result["download_way"]
        insert_record["JobID"] = ht_job[result["JobID"]]

        if "file_size" in result:
            insert_record["file_size"] = result["file_size"]
        if "file_format" in result:
            insert_record["file_format"] = result["file_format"]
        if "data_store_path" in result:
            paths = []
            for path in result["data_store_path"]:
                paths.append(path.replace("p300.icb.ac.cn:/mnt/p300/bdp/downloads", "/bdp-picb/downloads"))
            insert_record["data_store_path"] = paths

        if "NODE_path" in result:
            insert_record["NODE_path"] = result["NODE_path"]
        insert_record["NODE"] = result["NODE"]
        insert_record["to_node_date"] = ht_date2[result["JobID"]]
        collection_antnest.insert_one(insert_record)
        count += 1
        if count % 1000 == 0:
            print(str(count) + " records processed.")
    print("Finish processing assigned data.")

    # 处理还未分配的
    count = 0
    print("Start processing unassigned data.")
    for result in collection_raw.find({"JobID": {"$exists": 0}}):
        result2 = collection_antnest.find_one({"RunID": result["RunID"]}, {"_id": 1})
        if result2 == None:
            insert_record = {}
            if "BioprojectID" in result:
                insert_record["BioprojectID"] = result["BioprojectID"]
            if "Study" in result:
                insert_record["StudyID"] = result["Study"]
            insert_record["RunID"] = result["RunID"]
            data_download_url = "anonftp@ftp.ncbi.nlm.nih.gov:/sra/sra-instant/reads/ByRun/sra/" + result["RunID"][
                                                                                                   :3] + "/" + \
                                result[
                                    "RunID"][
                                :6] + "/" + \
                                result["RunID"] + "/" + result["RunID"] + ".sra"
            insert_record["data_download_url"] = data_download_url
            if "source" in result:
                insert_record["source"] = result["source"]
            if "backup" in result:
                insert_record["backup"] = result["backup"]
            if "type" in result:
                insert_record["type"] = result["type"]
            insert_record["data_create_date"] = datetime.today()

            insert_record["download_count"] = 0
            insert_record["NODE"] = "N"
            collection_antnest.insert_one(insert_record)
        else:
            if "type" in result:
                collection_antnest.update_one({"RunID": result["RunID"]}, {"$set": {"type": result["type"]}})
        count += 1
        if count % 1000 == 0:
            print(str(count) + " records processed.")

    print("Finish processing unassigned data.")

    CloseMongoDB(client_raw)
    CloseMongoDB(client_antnest)


def Test_Antnest():
    client_raw, collection_raw = InitMongoDB("localhost", 27017, "PD_Data", "rawdata")
    client_antnest, collection_antnest = InitMongoDB("localhost", 27017, "ant_nest", "datalist")

    for result in collection_antnest.find({}, {"RunID": 1}):
        result2 = collection_raw.find_one({"RunID": result["RunID"]}, {"data_store_path": 1})
        if "data_store_path" in result2:
            print(result["RunID"])
    CloseMongoDB(client_raw)
    CloseMongoDB(client_antnest)


if __name__ == '__main__':
    # TypeRawdata2Datalist("Gut")
    # TypeRawdata2Datalist("cancer")
    # TypeRawdata2Datalist("human and animal gut microbiome")
    # TypeRawdata2Datalist("gut microbiome")
    # TypeRawdata2Datalist("Host-associated:Human:Respiratory system:Pulmonary system:Lung ")
    # TypeRawdata2Datalist("hydrosphere")
    # NoTypeRawdata2Datalist()
    fire.Fire()
